import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
import pandas as pd
import os


def create_linear_model():
    model = nn.Sequential(
        nn.Linear(1, 32),
        nn.ReLU(),
        nn.Linear(32, 128),
        nn.ReLU(),
        nn.Linear(128, 64),
        nn.ReLU(),
        nn.Linear(64, 1))
    return model


def data():
    # 数据加载
    input_path = os.path.join('数据', '数据扩充.xlsx')
    output_path = os.path.join('数据', '数据.xlsx')

    dataset_new = pd.read_excel(input_path, sheet_name='整体数据')
    dataset = pd.read_excel(output_path, sheet_name='Sheet1')

    # 特征处理
    x = dataset_new.iloc[:, 0].values.reshape(-1, 1).astype(float)
    y = dataset.iloc[:, 3].values.reshape(-1, 1).astype(float)

    # 标准化参数计算
    # cv = 0.25
    # mean_x = np.mean(x)
    # std_x = mean_x * cv
    # mean_y = np.mean(y)
    # std_y = mean_y * cv

    mean_x = np.mean(x)
    std_x = np.std(x)
    mean_y = np.mean(y)
    std_y = np.std(y)

    # 特征标准化
    x_normalized = (x - mean_x) / std_x

    return x_normalized, mean_y, std_y


def test(model, test_loader, mean_y, std_y):
    model.eval()
    pre_outputs = []

    with torch.no_grad():
        for inputs in test_loader:  # 正确解包方式
            if isinstance(inputs, (list, tuple)):
                inputs = inputs[0]

            outputs = model(inputs)
            # 反标准化
            predictions = outputs.cpu().numpy() * std_y + mean_y
            pre_outputs.append(predictions)

    # 合并预测结果
    pre_outputs = np.concatenate(pre_outputs).flatten()

    # 写入文件
    input_path = os.path.join('数据', '数据扩充.xlsx')
    output_path = os.path.join('数据', 'DNN3_预测结果.xlsx')

    try:
        df = pd.read_excel(input_path, sheet_name='整体数据')

        # 校验数据长度
        if len(df) != len(pre_outputs):
            raise ValueError(f"数据长度不匹配！原始数据：{len(df)}，预测结果：{len(pre_outputs)}")

        # 添加预测列并保存
        df['模型预测值'] = pre_outputs
        df.to_excel(output_path, index=False)
        print(f"预测结果已成功保存至：{output_path}")

    except Exception as e:
        print(f"文件保存失败：{str(e)}")


if __name__ == "__main__":
    # 加载数据
    train_inputs, mean_y, std_y = data()

    # 创建数据集
    train_dataset = TensorDataset(torch.tensor(train_inputs, dtype=torch.float32))

    # 数据加载器
    BATCH_SIZE = 32
    train_loader = DataLoader(train_dataset,batch_size=BATCH_SIZE,shuffle=False)

    # 初始化并加载模型
    model = create_linear_model()
    try:
        model.load_state_dict(
            torch.load(
                os.path.join('模型', 'B_DNN3_model.pth'),
                map_location=torch.device('cpu')  # 确保设备兼容
            )
        )
    except Exception as e:
        print(f"模型加载失败：{str(e)}")
        exit(1)

    # 执行预测
    test(model, train_loader, mean_y, std_y)